package com.db.manage.datasource;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.lang.Nullable;

public class DynamicDataSource extends AbstractRoutingDataSource {
    private final Map<Object, Object> targetDataSources = new ConcurrentHashMap<>();
    private Object defaultTargetDataSource;

    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDataSource();
    }

    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) {
        this.targetDataSources.clear();
        this.targetDataSources.putAll(targetDataSources);
        super.setTargetDataSources(this.targetDataSources);
    }

    @Override
    public void setDefaultTargetDataSource(@Nullable Object defaultTargetDataSource) {
        this.defaultTargetDataSource = defaultTargetDataSource;
        super.setDefaultTargetDataSource(defaultTargetDataSource);
    }

    public Map<Object, Object> getTargetDataSources() {
        return targetDataSources;
    }

    public void setCurrentLookupKey(String key) {
        DataSourceContextHolder.setDataSource(key);
    }
} 